import cv2
import os


def template_matching(target_name="", template_name="", other_dir=""):

    nonsuffixal_name = os.path.splitext(target_name)[0]
    suffixal_name = os.path.splitext(target_name)[1]
    target = cv2.imread(target_name)
    template = cv2.imread(template_name)
    theight, twidth = template.shape[:2]
    result = cv2.matchTemplate(target, template, cv2.TM_SQDIFF_NORMED)
    cv2.normalize(result, result, 0, 1, cv2.NORM_MINMAX, -1)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)

    target_cropped = target[min_loc[1]:min_loc[1] + twidth,
                            min_loc[0]:min_loc[0] + theight]
    cv2.imwrite(nonsuffixal_name + "_cropped" + suffixal_name, target_cropped)

    cv2.rectangle(target, min_loc, (min_loc[0] + twidth, min_loc[1] + theight),
                  (0, 255, 0), 3)

    cv2.imwrite(nonsuffixal_name + "_marked" + suffixal_name, target)
    print("生成 " + nonsuffixal_name + "_cropped" + suffixal_name, end="\t")
    print(nonsuffixal_name + "_marked" + suffixal_name)

    if not other_dir == "":
        name_list = [
            n for n in os.listdir(other_dir)
            if os.path.isfile(os.path.join(other_dir, n))
        ]
        for name in name_list:
            img_path = os.path.join(other_dir, name)
            other_img = cv2.imread(img_path)

            cropped_other_img = other_img[min_loc[1]:min_loc[1] + twidth,
                                          min_loc[0]:min_loc[0] + theight]

            save_dir = os.path.join(other_dir, 'cropped')
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(save_dir, name)

            nonsuffixal_name = os.path.splitext(save_path)[0]
            suffixal_name = os.path.splitext(save_path)[1]
            cv2.imwrite(nonsuffixal_name + "_cropped" + suffixal_name,
                        cropped_other_img)

            print("生成 " + nonsuffixal_name + "_cropped" + suffixal_name)

            cv2.rectangle(other_img, min_loc,
                          (min_loc[0] + twidth, min_loc[1] + theight),
                          (0, 255, 0), 3)

            save_dir = os.path.join(other_dir, 'marked')
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(save_dir, name)

            nonsuffixal_name = os.path.splitext(save_path)[0]
            suffixal_name = os.path.splitext(save_path)[1]
            cv2.imwrite(nonsuffixal_name + "_marked" + suffixal_name, other_img)

            print("生成 " + nonsuffixal_name + "_marked" + suffixal_name)


if __name__ == "__main__":
    template_matching(target_name="target/8-014.png",
                      template_name="template/8-014.png",
                      other_dir="other_dir/ddan/ddan-300k-8")
